'''
Copyright (c) 2021 https://gitee.com/intellen
Intellen - Network is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2. 
You may obtain a copy of Mulan PSL v2 at:
            http://license.coscl.org.cn/MulanPSL2 
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.  
See the Mulan PSL v2 for more details. 

Url     : https://gitee.com/intellen/network
'''
from socketserver import ThreadingTCPServer, StreamRequestHandler
from Ucode import *
import threading

CONN_WORD = '123456'
CONN_PORT = 1105
CONT_WORD = 'kInGNmNOKdES6MCL'

TIME_OUT = 120

# 连接管理 - 储存连接实例
class DictBoxcon:
    '''Manage nodes'''
    def __init__(self):
        self.list = {}
    
    def add(self, con:Consiver, cnt):
        # 添加实例
        uid = con.get('uid')
        if uid == None:return False
        self.list[uid] = [con, cnt]
        return True
    
    def alive(self):
        for uid in self.list:
            con = self.list[uid]
            try:con[0].send(''.encode('utf-8'))
            except:
                con[1].finish()
                self.list.pop(uid)

    def listall(self):
        data = {}
        for uid, con in self.list.items():
            data[uid] = con[0].Idata
        return data
    
    def search(self, uid):
        con = self.list.get(uid)
        if con:return con[0]
    
    def remove(self, uid):
        con = self.search(uid)
        if con == None:return False
        try:con[1].finish()
        except:pass
        self.list.pop(uid)
        return True

conbox = DictBoxcon()

class CoreTree(StreamRequestHandler):

    def setup(self):
        self.request.settimeout(TIME_OUT)

        '''Confirm Connection'''
        # Init
        self.confirm = False
        self.uid = ''
        self.cons = Consiver(self.request, 'TCP')

        # Recv Data
        try:
            data = self.cons.recv().decode('utf-8')
            data = ast.literal_eval(data)
            if type(data) != dict:return
        except:return

        # Confirm and Save
        if data.get('_PWD') != CONN_WORD:return
        for key, values in data.items():
            if key[0] == '_':continue
            self.cons.save(key, values)
        self.uid = self.cons.get('uid')
        if not self.uid:self.cons.close()

        # 重复连接删除
        old_con = conbox.search(self.uid)
        if old_con:old_con.close()
        
        try:self.cons.send(self.uid.encode('utf-8'))
        except:return
        conbox.add(self.cons, self)
        self.confirm = True
        print('I:', self.uid)
    
    def trans(self, data):
        target = data.get('target')
        print('T:', target)
        size = int(data.get('size'))
        buff = int(data.get('buff', '1024'))
        # ------ Waiting to be fixed

        con = conbox.search(target)
        if not con:
            self.cons.send('Not'.encode('utf-8'))
            return

        # -------
        con.send(str([data.get('name','temp.cot'), size, buff]).encode('utf-8'))
        self.cons.send('Con'.encode('utf-8'))
        while size > 0:
            if size - 2048 > 0:
                temp = self.cons.conn.recv(2048)
            else:
                temp = self.cons.conn.recv(size)
            con.conn.send(temp)
            size -= len(temp)
        con.close()
    
    def execute(self, data):
        # Get command
        reda = ''
        cmd = data.pop('cmd')

        # Do something
        if cmd == 'lsc':
            conbox.alive()
            reda = conbox.listall()
        if cmd == 'trans':
            self.trans(data)
            return
        
        # Return result
        ret = {
            'sender':'_server',
            'time':time.time(),
            '_add':data.get('_add', ''),
            '_data':reda
        }
        try:self.cons.send(str(ret).encode('utf-8'))
        except:self.cons.close()
    
    def forward(self, data):
        con = conbox.search(data.pop('recver'))
        if not con:return
        met = data.pop('method', 'clt')
        data['time'] = time.time()
        if met == 'clt':
            data['sender'] = self.uid
        elif met == 'sev':
            data['sender'] = '_server'
        else:
            data['sender'] = self.uid
        try:con.send(str(data).encode('utf-8'))
        except:con.close()
    
    def handle(self):
        if not self.confirm:return
        while True:
            try:
                data = self.cons.recv().decode('utf-8')
                data = ast.literal_eval(data)
            except:return
            dtype = data.pop('_type')

            if dtype == 'non':continue

            # execute
            if dtype == 'cmd':
                self.execute(data)
            # Trans data
            if dtype == 'for':
                self.forward(data)

    def finish(self):
        conbox.remove(self.uid)
        print('O:', self.uid)


def start():
    print('Start on port', CONN_PORT)
    addr = ('0.0.0.0',CONN_PORT)
    server = ThreadingTCPServer(addr, CoreTree)
    server.serve_forever()


if __name__ == "__main__":
    start()